import os
import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk
from moviepy import VideoFileClip
import threading

# 定义全局变量用于统计成功和失败的数量
success_count = 0
failure_count = 0

def extract_first_frame(video_path, output_folder, video_folder, total_files, index, success_label, failure_label, total_label):
    # 获取视频文件的相对路径
    relative_path = os.path.relpath(video_path, video_folder)
    # 获取图像保存路径，注意添加 .jpg 扩展名
    output_image_path = os.path.join(output_folder, os.path.splitext(relative_path)[0] + '.jpg')

    # 确保输出文件夹存在
    os.makedirs(os.path.dirname(output_image_path), exist_ok=True)

    # 检查文件是否存在
    if os.path.exists(output_image_path):
        # 弹出确认窗口
        if messagebox.askyesno("文件已存在", f"{output_image_path} 已存在，是否要替换?"):
            pass  # 用户选择替换，继续执行
        else:
            return  # 用户选择不替换，退出此函数

    try:
        # 提取第一帧
        clip = VideoFileClip(video_path)
        clip.save_frame(output_image_path, t=0)
        clip.close()
        success_count += 1
        success_label.config(text=f"成功: {success_count}")  # 更新成功数标签
    except Exception as e:
        failure_count += 1
        failure_label.config(text=f"失败: {failure_count}")  # 更新失败数标签
        messagebox.showerror("错误", f"提取 {video_path} 的第一帧时出错: {str(e)}")

    # 更新总数标签
    total_label.config(text=f"总数: {total_files}")

    

    # 更新正在进行提取的视频标签
    current_video_label.config(text=f"正在提取: {os.path.basename(video_path)}")

def batch_extract_frames(video_folder, output_folder,  success_label, failure_label, total_label):
    video_files = []
    failure_count = 0
    total_files = 0  # 提前定义并初始化为0

    for root, _, files in os.walk(video_folder):
        for video_file in files:
            if video_file.endswith(('.mp4', '.avi', '.mkv', '.wmv', '.flv', '.MP4')):
                video_path = os.path.join(root, video_file)
                video_files.append(video_path)

                # 获取视频文件的相对路径
                relative_path = os.path.relpath(video_path, video_folder)
                # 获取图像保存路径，注意添加 .jpg 扩展名
                output_image_path = os.path.join(output_folder, os.path.splitext(relative_path)[0] + '.jpg')

                # 确保输出文件夹存在
                os.makedirs(os.path.dirname(output_image_path), exist_ok=True)

                # 检查文件是否存在
                if os.path.exists(output_image_path):
                    # 弹出确认窗口
                    if messagebox.askyesno("文件已存在", f"{output_image_path} 已存在，是否要替换?"):
                        pass  # 用户选择替换，继续执行
                    else:
                        continue  # 用户选择不替换，跳过此文件

                try:
                    # 提取第一帧
                    clip = VideoFileClip(video_path)
                    clip.save_frame(output_image_path, t=0)
                    clip.close()

                    # 检查图片大小
                    if os.path.getsize(output_image_path) > 150 * 1024:  # 150kb
                        # 如果图片大小超过限制，进行压缩
                        from PIL import Image
                        with Image.open(output_image_path) as img:
                            img.save(output_image_path, optimize=True, quality=85)  # 调整quality参数以达到合适的压缩效果

                    success_count += 1
                    success_label.config(text=f"成功: {success_count}")  # 更新成功数标签
                except Exception as e:
                    failure_count += 1
                    failure_label.config(text=f"失败: {failure_count}")  # 更新失败数标签
                    messagebox.showerror("错误", f"提取 {video_path} 的第一帧时出错: {str(e)}")

                # 更新总数标签
                total_files = len(video_files)  # 在循环中更新 total_files
                total_label.config(text=f"总数: {total_files}")

    def worker():
        global success_count, failure_count
        success_count = 0
        failure_count = 0
        for index, video_path in enumerate(video_files):  # 在这里定义 index
            extract_first_frame(video_path, output_folder, video_folder,total_files, index, success_label, failure_label, total_label)

        messagebox.showinfo("完成", f"所有视频的第一帧已提取完成\n成功: {success_count}\n失败: {failure_count}\n总数: {total_files}")

    # 创建并启动线程
    thread = threading.Thread(target=worker)
    thread.start()
def select_video_folder():
    folder = filedialog.askdirectory()
    if folder:
        video_folder_var.set(folder)

def select_output_folder():
    folder = filedialog.askdirectory()
    if folder:
        output_folder_var.set(folder)

def start_extraction():
    video_folder = video_folder_var.get()
    output_folder = output_folder_var.get()

    if not video_folder or not output_folder:
        messagebox.showwarning("警告", "请先选择视频文件夹和输出文件夹")
        return

    
    batch_extract_frames(video_folder, output_folder, success_label, failure_label, total_label)

# 创建主窗口
root = tk.Tk()
root.title("视频帧提取工具")

# 视频文件夹选择
video_folder_var = tk.StringVar()
tk.Label(root, text="视频文件夹:").grid(row=0, column=0, padx=5, pady=5)
tk.Entry(root, textvariable=video_folder_var, width=50).grid(row=0, column=1, padx=5, pady=5)
tk.Button(root, text="选择文件夹", command=select_video_folder).grid(row=0, column=2, padx=5, pady=5)

# 输出文件夹选择
output_folder_var = tk.StringVar()
tk.Label(root, text="输出文件夹:").grid(row=1, column=0, padx=5, pady=5)
tk.Entry(root, textvariable=output_folder_var, width=50).grid(row=1, column=1, padx=5, pady=5)
tk.Button(root, text="选择文件夹", command=select_output_folder).grid(row=1, column=2, padx=5, pady=5)


# 创建正在进行提取的视频标签
current_video_label = tk.Label(root, text="正在提取: ")
current_video_label.grid(row=3, column=0, columnspan=3, padx=5, pady=5)

# 创建成功数、失败数和总数标签
success_label = tk.Label(root, text="成功: 0")
success_label.grid(row=4, column=0, padx=5, pady=5)
failure_label = tk.Label(root, text="失败: 0")
failure_label.grid(row=4, column=1, padx=5, pady=5)
total_label = tk.Label(root, text="总数: 0")
total_label.grid(row=4, column=2, padx=5, pady=5)

# 开始提取按钮
tk.Button(root, text="开始提取", command=start_extraction).grid(row=5, column=1, pady=10)

# 运行主循环
root.mainloop()
